import os

import cv2
import numpy as np
from skimage.metrics import structural_similarity

from tools.flip import compute_flip
from tools.utils import chw_to_hwc, index2color, get_magma_map


def abs_error(img1, img2, diff_image_path):
    if os.path.exists(diff_image_path):
        return
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    diff = np.abs(img1 - img2)
    diff = diff.astype(np.uint8)
    diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(diff_image_path, diff)


def squared_error(img1, img2, diff_image_path):
    if os.path.exists(diff_image_path):
        return
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    diff = np.square(img1 - img2)
    diff = diff.astype(np.uint8)
    diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(diff_image_path, diff)


def psnr(img1, img2, diff_image_path):
    if os.path.exists(diff_image_path):
        return
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)
    diff = np.square(img1 - img2)
    diff = 10 * np.log10((255 ** 2) / diff)
    diff = diff.astype(np.uint8)
    diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(diff_image_path, diff)


def ssim(img1, img2, diff_image_path, gaussian_weights=False):
    if os.path.exists(diff_image_path):
        return
    d, a = structural_similarity(img1, img2, multichannel=True, gaussian_weights=gaussian_weights, full=True)
    a = 1 - a
    diff = (a * 255).astype(np.uint8)
    diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(diff_image_path, diff)


def flip(img1, img2, diff_image_path, monitor_distance=0.7, monitor_width=0.7, monitor_resolution_x=3840,
         use_color_map=True):
    if os.path.exists(diff_image_path):
        return
    # Compute number of pixels per degree of visual angle
    pixels_per_degree = monitor_distance * (monitor_resolution_x / monitor_width) * (np.pi / 180)
    # Compute FLIP map
    deltaE = compute_flip(img1, img2, pixels_per_degree)
    # Save error map
    index_map = np.floor(255.0 * deltaE.squeeze(0))

    if use_color_map:
        result = chw_to_hwc(index2color(index_map, get_magma_map()))
    else:
        result = index_map / 255.0
        # save_image(diff_image_path, result)
    result = (np.clip(result, 0.0, 1.0) * 255.0).astype(np.uint8)
    result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(diff_image_path, result)


def get_real_fake_pairs(args):
    real_fake_pairs = []
    lines = open(os.path.join(args.dst_dir, 'real_fake_pairs.dat')).readlines()
    for line in lines:
        line_split = line.split(',')
        real_fake_pairs.append((line_split[0], line_split[1].strip()))

    return real_fake_pairs
